
import os
import sys
import time
import traceback
import multiprocessing as mp

import rpyc
from rpyc.utils import server

import dexen
from dexen.system import util, node
from dexen.system.base import service, MyManager, task
from dexen.system.base.spec import DexenSpec
#from dexen.system.base.exceptions import ScriptException, SlaveException
from dexen.system.node.slave import comm

#===============================================================================
# Globals
#===============================================================================
MyManager.register("SlaveComm", comm.SlaveComm)

#===============================================================================
# Starting
#===============================================================================
def start(_mastername, _slavename, _nodename):
    global mastername, slavename, slavedir, nodename, manager
    mastername = _mastername
    slavename = _slavename
    nodename = _nodename
    manager = MyManager()
    manager.start()
    
    slavedir = os.path.join(node.TEMP_DIR, slavename)
    if not os.path.exists(slavedir):
        os.mkdir(slavedir)
    sys.path.append(slavedir)
    
    SlaveService.ALIASES = [slavename]
    svr = server.ThreadedServer(SlaveService,
                                protocol_config=service.default_config)
    svr.start()
        
#===============================================================================
# Service Class
#===============================================================================
class SlaveService(dexen.Service):
    def on_connect(self):
        pass
    
    def on_disconnect(self):
        pass

    def run(self, spec):
        #print "I am in slave these are the tasks", tasks
        assert isinstance (spec, DexenSpec) 
        spec = rpyc.classic.obtain(spec)
        print "The slave dir is", slavedir
        util.unzip(spec.schema_zip, slavedir)
        self.comm = manager.SlaveComm(slavename, slavedir)
        spec.slavedir = slavedir
        print "The execute_tasks process is being initiated"
        self.proc = mp.Process(target=execute_tasks, 
                               args=(spec, self.comm))
        self.proc.start()
        
    def put(self, task_name, input):
        try:
            input = rpyc.classic.obtain(input)
            unit = task.InputUnit(task_name, input)
            self.comm.put_input(unit)
        except:
            traceback.print_exc()
    
    def get(self):
        return self.comm.get_outputs()
    
    def get_rank(self):
        return self.comm.get_rank()
    
    def available(self):
        if self.comm.get_exc():
            raise self.comm.get_exc()
        return self.comm.can_put()
    
#===============================================================================
# TASK EXECUTOR
#===============================================================================
def execute_tasks(spec, s_comm):
    from StringIO import StringIO
    msg = StringIO()
    assert isinstance(spec, DexenSpec)
    
    sys.path.append(spec.slavedir)
    os.chdir(spec.slavedir)
    
    spec.objs = {} #empties the spec.objs
    
    while True:
        try:
            while True:
                i_unit = s_comm.get_input()
                if i_unit:
                    break
                time.sleep(1)        
            task_name = i_unit.task_name
            
            assert isinstance(i_unit, task.InputUnit) 
            print "executing", task_name
            
            obj = spec.get_task_obj(task_name)
            
            try:
                inds = obj.__execute__(i_unit.input)
            except:
                traceback.print_exc(file=msg)
                s_comm.crash(comm.SCRIPT_CRASHED, msg)
                break    
            else:
                o_unit = task.OutputUnit(task_name, inds)
                s_comm.put_output(o_unit)    
        except Exception:
            traceback.print_exc(file=msg)
            s_comm.crash(comm.SLAVE_CRASHED, msg)
            print "Slave has an exception I am exiting"
            break
 
#===============================================================================
# TESTS
#===============================================================================
def test_suite():
    pass

if __name__ == '__main__':
    test_suite()
    

